using UnityEngine;
using UnityEngine.Rendering;
using UnityEngine.Experimental.Rendering;
using System.Collections.Generic;
using System;
    
public class HRenderPipelineInstance : RenderPipeline
{
    private CommandBuffer m_gcb;
    private CommandBuffer m_vcb;
    private CommandBuffer m_rcb;
    private CommandBuffer m_scb;
    private int width;
    private int height; 

    RenderTargetIdentifier defaultRT_identifier;
    RenderTargetIdentifier vbuffer_identifier;

    private bool useRasterVBuffer = false;
    private bool useVBufferDebug = true;
    private bool debugEnd = false;
    //visibility pass
    private RenderTexture m_VisibilityBuffer;
    private ComputeShader unpackCompute;
    private int kernelID_Unpack;
    private ComputeBuffer unpackedIDBuffer;
    private const int INSTANCE_ID_MASK = 1023;
    private const int INSTANCE_ID_BITS = 10;
    private Texture2D idBuffer;
    private uint[] id_list;
    private HashSet<uint>[] Inst_Prims;

    private List<Mesh> mesh_list; // save original meshes
    private GameObject[] totalObjects;
    private List<GameObject> obj_list; // save renderable GameObjects
    private Material[] materials;
    private int objCount;
    private List<int> id2mat; // InstanceID -> Material ID

    // rasterizing pass
    private RenderTexture m_BlitBuffer;
    private ComputeBuffer m_vertexBuffer;
    private ComputeBuffer m_indexBuffer;
    private ComputeBuffer m_instID2Index;
    private ComputeBuffer m_uvBuffer;
    private ComputeBuffer m_transformBuffer;
    private Texture2DArray m_textures;
    ComputeShader rasterCompute;
    private int kernelID_Raster;

    private int cnt = 0;
    /*
        public static Mesh FullScreenMesh_Game
        {
            get
            {
                if (FullScreenMeshOrigin_Game != null) {
                    return FullScreenMeshOrigin_Game;
                }

                FullScreenMeshOrigin_Game = new Mesh { name = "FullScreen Mesh" };

                FullScreenMeshOrigin_Game.vertices = new Vector3[] {
                    new Vector3(-1, -1, 0),
                    new Vector3(-1,  1, 0),
                    new Vector3( 1,  1, 0),
                    new Vector3( 1, -1, 0)
                };

                FullScreenMeshOrigin_Game.uv = new Vector2[] {
                    new Vector2(0, 0),
                    new Vector2(0, 1),
                    new Vector2(1, 1),
                    new Vector2(1, 0)
                };

                FullScreenMeshOrigin_Game.SetIndices(new int[] { 0, 1, 2, 3 }, MeshTopology.Quads, 0, false);
                FullScreenMeshOrigin_Game.UploadMeshData(false);
                return FullScreenMeshOrigin_Game;
            }
        }
    */
    struct WorkItem
    {
        public uint pID;
        public int instID;
    };
    public struct uint4
    {
        public uint x;
        public uint y;
        public uint z;
        public uint w;
        public uint4 (uint _x, uint _y, uint _z, uint _w)
            {
            x = _x;
            y = _y;
            z = _z;
            w = _w;
            }
    }

    public HRenderPipelineInstance(ComputeShader unpackCompute, ComputeShader rasterCompute) 
    {
        m_gcb = new CommandBuffer();
        m_gcb.name = "game command buffer";
        m_vcb = new CommandBuffer();
        m_vcb.name = "vbuffer command buffer";
        m_rcb = new CommandBuffer();
        m_rcb.name = "rasterize command buffer";
        m_scb = new CommandBuffer();
        m_scb.name = "scene command buffer";

        mesh_list = new List<Mesh>();
        obj_list = new List<GameObject>();
        id2mat = new List<int>();

        this.unpackCompute = unpackCompute;
        kernelID_Unpack = unpackCompute.FindKernel("parseIDs");

        this.rasterCompute = rasterCompute;
        kernelID_Raster = rasterCompute.FindKernel("rasterPass");

        materials = new Material[3];
        materials[0] = Resources.Load("H1_mat", typeof(Material)) as Material;
        materials[1] = Resources.Load("H2_mat", typeof(Material)) as Material;
        materials[2] = Resources.Load("H3_mat", typeof(Material)) as Material;
    }
    public void countRenderObjs()
    {
        totalObjects = UnityEngine.Object.FindObjectsOfType(typeof(GameObject)) as GameObject[];
        objCount = 0;

        foreach (var obj in totalObjects)
        {
            if (obj.GetComponent<Renderer>() != null)
            { 
                Mesh obj_mesh = UnityEngine.Object.Instantiate(obj.GetComponent<MeshFilter>().mesh);
                mesh_list.Add(obj_mesh);
                obj_list.Add(obj);
                objCount++;
            }   
        }

        Inst_Prims = new HashSet<uint>[objCount];

        for (int i = 0; i < objCount; i++)
        {
            Inst_Prims[i] = new HashSet<uint>();
        }
        /**/
        int k = 0;
        uint vert_off = 0;
        uint idx_off = 0;
        uint uv_off = 0;
        List<Vector3> vertex_cpu = new List<Vector3>();
        List<Vector2> uv_cpu = new List<Vector2>();
        List<int> idx_cpu = new List<int>();
        List<uint4> inst2idx_cpu = new List<uint4>();
        List<uint> id2mat_cpu = new List<uint>();
        List<Matrix4x4> matrix_cpu = new List<Matrix4x4>();
        foreach (var obj in obj_list)
        {
            Material mat = obj.GetComponent<Renderer>().material;
            int mat_i = 4;
            if (mat.name == "H1_mat (Instance)")
                mat_i = 0;
            else if (mat.name == "H2_mat (Instance)")
                mat_i = 1;
            else if (mat.name == "H3_mat (Instance)")
                mat_i = 2;

            uint4 idxs = new uint4(vert_off, idx_off, uv_off, (uint)mat_i);
            inst2idx_cpu.Add(idxs);

            Mesh mesh = obj.GetComponent<MeshFilter>().mesh;
            vertex_cpu.AddRange(mesh.vertices);
            uv_cpu.AddRange(mesh.uv);
            idx_cpu.AddRange(mesh.triangles);
            matrix_cpu.Add(obj.transform.localToWorldMatrix);

            idx_off += (uint)mesh.triangles.Length;
            uv_off += (uint)mesh.uv.Length;
            vert_off += (uint)mesh.vertices.Length;
            Debug.Log("triangles : " + mesh.triangles.Length + " uvs : " 
                    + mesh.uv.Length + " verts : " + mesh.vertices.Length + " normals : " + mesh.normals.Length);
        }
        m_vertexBuffer = new ComputeBuffer(vertex_cpu.Count, sizeof(float)*3);
        m_uvBuffer = new ComputeBuffer(uv_cpu.Count, sizeof(float) * 2);
        m_indexBuffer = new ComputeBuffer(idx_cpu.Count, sizeof(int));
        m_instID2Index = new ComputeBuffer(objCount, sizeof(uint)*4);
        m_transformBuffer = new ComputeBuffer(objCount, sizeof(float) * 16);
        m_vertexBuffer.SetData(vertex_cpu);
        m_uvBuffer.SetData(uv_cpu);
        m_indexBuffer.SetData(idx_cpu);
        m_instID2Index.SetData(inst2idx_cpu);
        m_transformBuffer.SetData(matrix_cpu);
        //material
    }
    
    private void reGenerateMeshes()
    {
        for(int i=0;i<objCount;i++)
        {
            obj_list[i].GetComponent<MeshFilter>().mesh = mesh_list[i];
        }
    }

    private void buildVBuffer(ref ScriptableRenderContext context)
    {
        if (id_list == null)
        {
            id_list = new uint[width * height];
        }
        m_VisibilityBuffer = RenderTexture.GetTemporary(width, height, 16, GraphicsFormat.R32_UInt);
        m_VisibilityBuffer.name = "visibility_buffer";

        vbuffer_identifier = new RenderTargetIdentifier(m_VisibilityBuffer);
        m_vcb.SetRenderTarget(vbuffer_identifier);
        m_vcb.SetViewport(new Rect(0, 0, width, height));
        m_vcb.ClearRenderTarget(true, true, Color.white);

        for (int id=0; id < objCount; id++)
        {
            m_vcb.SetGlobalInt(Shader.PropertyToID("InstanceID"), id);
            Renderer renderer = obj_list[id].GetComponent<Renderer>();
            var mat = renderer.sharedMaterial;
            m_vcb.DrawRenderer(renderer, mat, 0, 1); //HUnlit pass 1 (vbuffer)      
        }
        context.ExecuteCommandBuffer(m_vcb);    
        m_vcb.Clear();
    }

    private void rasterize(ref ScriptableRenderContext context, in Camera cam)
    {
        //build vertex buffer
        m_BlitBuffer = RenderTexture.GetTemporary(width, height, 0, RenderTextureFormat.ARGB32);
        m_BlitBuffer.name = "blit buffer";
        m_BlitBuffer.enableRandomWrite = true;

        Matrix4x4 world2Screen = cam.projectionMatrix * cam.worldToCameraMatrix;
        Matrix4x4 screen2World = world2Screen.inverse;

        rasterCompute.SetInt("height", cam.pixelHeight);
        rasterCompute.SetInt("width", cam.pixelWidth);
        rasterCompute.SetFloat("nearPlane", cam.nearClipPlane);
        rasterCompute.SetMatrix("S2W", screen2World);

        rasterCompute.SetTexture(kernelID_Raster, "Result", m_BlitBuffer);
        rasterCompute.SetTexture(kernelID_Raster, "packedIDBuffer", m_VisibilityBuffer);

        rasterCompute.SetBuffer(kernelID_Raster, "vertexBuffer", m_vertexBuffer);
        rasterCompute.SetBuffer(kernelID_Raster, "indexBuffer", m_indexBuffer);
        rasterCompute.SetBuffer(kernelID_Raster, "instID2Index", m_instID2Index);
        rasterCompute.SetBuffer(kernelID_Raster, "uvBuffer", m_uvBuffer);
        rasterCompute.SetBuffer(kernelID_Raster, "transformBuffer", m_transformBuffer);

        rasterCompute.SetTexture(kernelID_Raster, "_Tex0", materials[0].mainTexture);
        rasterCompute.SetTexture(kernelID_Raster, "_Tex1", materials[1].mainTexture);
        rasterCompute.SetTexture(kernelID_Raster, "_Tex2", materials[2].mainTexture);

        m_rcb.DispatchCompute(rasterCompute, kernelID_Raster, width / 16 + 1, height / 16 + 1, 1);
        context.ExecuteCommandBuffer(m_rcb);
        m_rcb.Clear();
    }

    private void buildVBufferNow()
    {
        if (id_list == null)
        {
            id_list = new uint[width * height];
        }
        m_VisibilityBuffer = RenderTexture.GetTemporary(width, height, 16, GraphicsFormat.R32_UInt);
        m_VisibilityBuffer.name = "visibility_buffer";

        Graphics.SetRenderTarget(m_VisibilityBuffer);
        GL.Clear(true, true, Color.white);

        for (int id = 0; id < objCount; id++)
        {
            Shader.SetGlobalInt(Shader.PropertyToID("InstanceID"), id);
            Material mat = obj_list[id].GetComponent<Renderer>().sharedMaterial;
            Mesh mesh = obj_list[id].GetComponent<MeshFilter>().mesh;
            mat.SetPass(1);
            Graphics.DrawMeshNow(mesh, obj_list[id].transform.localToWorldMatrix);
        }
    }
    private void buildMeshesNow()
    {
        cnt++;

        if (unpackedIDBuffer == null)
        {
            unpackedIDBuffer = new ComputeBuffer(width * height, sizeof(uint));
        }
        unpackCompute.SetBuffer(kernelID_Unpack, "unpackedIDBuffer", unpackedIDBuffer);
        unpackCompute.SetTexture(kernelID_Unpack, "packedIDBuffer", m_VisibilityBuffer);
        unpackCompute.SetInt("width", width);
        unpackCompute.SetInt("height", height);
        unpackCompute.Dispatch(kernelID_Unpack, width / 16 + 1, height / 16 + 1, 1);
        unpackedIDBuffer.GetData(id_list);

        for (int i = 0; i < objCount; i++)
        {
            Inst_Prims[i].Clear();
        }
        for(int i = 0; i < id_list.Length; i++)
        {
            uint packID = id_list[i];
            if (packID == 0xFFFFFFFF)
                continue;
            uint pID    = packID >> INSTANCE_ID_BITS;
            uint instID = packID & INSTANCE_ID_MASK;
            if (!Inst_Prims[instID].Contains(pID))
            {
                Inst_Prims[instID].Add(pID);
            }
        }
        
        List<int> tris = new List<int>();
        for (int i = 0; i < objCount; i++)
        {
            if(cnt == 2)
            {
                debugEnd = true;
                int[] obj_tris = mesh_list[i].GetTriangles(0);
                foreach (uint pID in Inst_Prims[i])
                {
                    tris.Add(obj_tris[pID * 3]);
                    tris.Add(obj_tris[pID * 3 + 1]);
                    tris.Add(obj_tris[pID * 3 + 2]);
                }
                Mesh obj_mesh = obj_list[i].GetComponent<MeshFilter>().mesh;
                obj_mesh.SetIndices(tris, MeshTopology.Triangles, 0);
                obj_mesh.UploadMeshData(true);
                tris.Clear();
            }

        }
        /**/
    }

    private void RenderGame(ref ScriptableRenderContext context, in Camera camera)
    {
        width = camera.pixelWidth;
        height = camera.pixelHeight;
        context.SetupCameraProperties(camera);
     
        defaultRT_identifier = new RenderTargetIdentifier(camera.activeTexture);

        if(useRasterVBuffer)
        {
            buildVBuffer(ref context);
            rasterize(ref context, camera);
            m_gcb.Blit(m_BlitBuffer, camera.targetTexture);
            context.ExecuteCommandBuffer(m_gcb);
            m_gcb.Clear();
            if(m_BlitBuffer)
            {
                RenderTexture.ReleaseTemporary(m_BlitBuffer);
            }
        }
        else if(useVBufferDebug)
        {
            camera.TryGetCullingParameters(out var cullingParameters);
            var cullingResults = context.Cull(ref cullingParameters);

            //
            //if(!debugEnd)
            if(cnt < 2)
            {
                buildVBuffer(ref context);
                //buildVBufferNow();
                buildMeshesNow();
            }

            /*
            var visibleLights = cullingResults.visibleLights;
            foreach (var light in visibleLights)
            {
                Vector4 pos = light.localToWorldMatrix.GetColumn(2);
                if (light.lightType == LightType.Directional)// main light
                {
                    Vector4 LightDirection = new Vector4(-pos.x, -pos.y, -pos.z, 0);
                    Color LightColor = light.finalColor;

                    m_gcb.SetGlobalVector(Shader.PropertyToID("H_MainLightDirection"), LightDirection);
                    m_gcb.SetGlobalColor(Shader.PropertyToID("H_MainLightColor"), LightColor);
                    context.ExecuteCommandBuffer(m_gcb);
                    m_gcb.Clear();
                }
            }
            */

            m_gcb.SetRenderTarget(defaultRT_identifier);
            m_gcb.ClearRenderTarget(true, true, Color.black);
            var sortSettings = new SortingSettings(camera);
            sortSettings.criteria = SortingCriteria.CommonOpaque;
            DrawingSettings drawSettings = new DrawingSettings(new ShaderTagId("HUnlit"), sortSettings);
            var filterSettings = new FilteringSettings(RenderQueueRange.opaque);
            context.ExecuteCommandBuffer(m_gcb);
            m_gcb.Clear();
            context.DrawRenderers(cullingResults, ref drawSettings, ref filterSettings);
            context.DrawSkybox(camera);
        }   
        else
        {
            camera.TryGetCullingParameters(out var cullingParameters);
            var cullingResults = context.Cull(ref cullingParameters);

            var visibleLights = cullingResults.visibleLights;
            foreach (var light in visibleLights)
            {
                Vector4 pos = light.localToWorldMatrix.GetColumn(2);
                if (light.lightType == LightType.Directional)// main light
                {
                    Vector4 LightDirection = new Vector4(-pos.x, -pos.y, -pos.z, 0);
                    Color LightColor = light.finalColor;

                    m_gcb.SetGlobalVector(Shader.PropertyToID("H_MainLightDirection"), LightDirection);
                    m_gcb.SetGlobalColor(Shader.PropertyToID("H_MainLightColor"), LightColor);
                    context.ExecuteCommandBuffer(m_gcb);
                    m_gcb.Clear();
                }
            }

            m_gcb.SetRenderTarget(defaultRT_identifier);
            m_gcb.ClearRenderTarget(true, true, Color.black);
            var sortSettings = new SortingSettings(camera);
            sortSettings.criteria = SortingCriteria.CommonOpaque;
            DrawingSettings drawSettings = new DrawingSettings(new ShaderTagId("HUnlit"), sortSettings);
            var filterSettings = new FilteringSettings(RenderQueueRange.opaque);
            context.ExecuteCommandBuffer(m_gcb);
            m_gcb.Clear();
            context.DrawRenderers(cullingResults, ref drawSettings, ref filterSettings); 
            context.DrawSkybox(camera);
        }
       
        context.Submit();


        if(m_VisibilityBuffer)
        {
            RenderTexture.ReleaseTemporary(m_VisibilityBuffer);
        }
        
    }

    private void RenderScene(ref ScriptableRenderContext context, in Camera camera)
    {
        context.SetupCameraProperties(camera);

        camera.TryGetCullingParameters(out var cullingParameters);
        var cullingResults = context.Cull(ref cullingParameters);

        var visibleLights = cullingResults.visibleLights;
        foreach (var light in visibleLights)
        {
            Vector4 pos = light.localToWorldMatrix.GetColumn(2);
            if (light.lightType == LightType.Directional)// main light
            {
                Vector4 LightDirection = new Vector4(-pos.x, -pos.y, -pos.z, 0);
                Color LightColor = light.finalColor;

                m_scb.SetGlobalVector(Shader.PropertyToID("H_MainLightDirection"), LightDirection);
                m_scb.SetGlobalColor(Shader.PropertyToID("H_MainLightColor"), LightColor);
            }
        }
        context.ExecuteCommandBuffer(m_gcb);
        m_scb.Clear();
        //opaque
        var sortSettings = new SortingSettings(camera);
        sortSettings.criteria = SortingCriteria.CommonOpaque;
        DrawingSettings drawSettings = new DrawingSettings(new ShaderTagId("HUnlit"), sortSettings);
        var filterSettings = new FilteringSettings(RenderQueueRange.opaque);

        context.DrawRenderers(cullingResults, ref drawSettings, ref filterSettings);
        context.DrawSkybox(camera);
        context.DrawGizmos(camera, GizmoSubset.PostImageEffects);
        context.Submit();
    }

    protected override void Render (ScriptableRenderContext context, Camera[] cameras) 
    {
        if (totalObjects == null)
        {  
            countRenderObjs();
        }
       
        // Iterate over all Cameras
        foreach (Camera camera in cameras)
        {        
            //var flags = camera.clearFlags;
            //m_gcb.ClearRenderTarget((flags & CameraClearFlags.Depth) != 0, (flags & CameraClearFlags) != 0, camera.backgroundColor);
           if(camera.cameraType == CameraType.Game)
            {
                RenderGame(ref context, camera);
            }
           else if(camera.cameraType == CameraType.SceneView)
            {
                RenderScene(ref context, camera);
            }
            
        }

    }

}